{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import sys\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('../')\n",
    "os.chdir('../experiments/siampose_ct/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.config_helper import load_config\n",
    "from datasets.siam_pose_ct_dataset import DataSets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self):\n",
    "        self.config = './config.json'\n",
    "        self.arch = 'custom'\n",
    "        self.epochs = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Args()\n",
    "cfg = load_config(args)\n",
    "train_set = DataSets(cfg['train_datasets'], cfg['anchors'], args.epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[[125.224625, 117.224625, 109.224625, ...,  88.224625,\n",
       "           81.224625,  76.224625],\n",
       "         [ 60.22463 ,  66.224625,  64.224625, ..., 104.224625,\n",
       "          113.224625, 115.224625],\n",
       "         [ 39.22463 ,  32.22463 ,  40.22463 , ..., 146.22462 ,\n",
       "          159.22462 , 157.22462 ],\n",
       "         ...,\n",
       "         [177.22462 , 175.22462 , 174.22462 , ..., 137.22462 ,\n",
       "          141.22462 , 141.22462 ],\n",
       "         [177.22462 , 176.22462 , 176.22462 , ..., 150.22462 ,\n",
       "          150.22462 , 149.22462 ],\n",
       "         [176.22462 , 175.22462 , 176.22462 , ..., 161.22462 ,\n",
       "          160.22462 , 158.22462 ]],\n",
       " \n",
       "        [[136.93405 , 122.93405 , 110.93405 , ...,  85.93405 ,\n",
       "           77.93405 ,  70.93405 ],\n",
       "         [ 74.93405 ,  74.93405 ,  69.93405 , ...,  95.93405 ,\n",
       "          101.93405 , 103.93405 ],\n",
       "         [ 57.93405 ,  46.93405 ,  51.93405 , ..., 125.93405 ,\n",
       "          132.93405 , 131.93405 ],\n",
       "         ...,\n",
       "         [162.93405 , 160.93405 , 159.93405 , ..., 124.93405 ,\n",
       "          128.93405 , 128.93405 ],\n",
       "         [162.93405 , 161.93405 , 161.93405 , ..., 137.93405 ,\n",
       "          137.93405 , 136.93405 ],\n",
       "         [161.93405 , 160.93405 , 161.93405 , ..., 146.93405 ,\n",
       "          145.93405 , 144.93405 ]],\n",
       " \n",
       "        [[209.17618 , 200.17618 , 190.17618 , ...,  89.17617 ,\n",
       "           84.17617 ,  82.17617 ],\n",
       "         [135.17618 , 137.17618 , 133.17618 , ...,  94.17617 ,\n",
       "          101.17617 , 106.17617 ],\n",
       "         [105.17617 ,  92.17617 ,  96.17617 , ..., 116.17617 ,\n",
       "          123.17617 , 123.17617 ],\n",
       "         ...,\n",
       "         [182.17618 , 180.17618 , 179.17618 , ..., 141.17618 ,\n",
       "          143.17618 , 144.17618 ],\n",
       "         [182.17618 , 181.17618 , 181.17618 , ..., 154.17618 ,\n",
       "          154.17618 , 152.17618 ],\n",
       "         [181.17618 , 180.17618 , 181.17618 , ..., 164.17618 ,\n",
       "          163.17618 , 160.17618 ]]], dtype=float32),\n",
       " array([[[ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032],\n",
       "         [ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032],\n",
       "         [ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032],\n",
       "         ...,\n",
       "         [ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032],\n",
       "         [ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032],\n",
       "         [ 91.93032,  91.93032,  91.93032, ...,  91.93032,  91.93032,\n",
       "           91.93032]],\n",
       " \n",
       "        [[107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746],\n",
       "         [107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746],\n",
       "         [107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746],\n",
       "         ...,\n",
       "         [107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746],\n",
       "         [107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746],\n",
       "         [107.14746, 107.14746, 107.14746, ..., 107.14746, 107.14746,\n",
       "          107.14746]],\n",
       " \n",
       "        [[114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821],\n",
       "         [114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821],\n",
       "         [114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821],\n",
       "         ...,\n",
       "         [114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821],\n",
       "         [114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821],\n",
       "         [114.09821, 114.09821, 114.09821, ..., 114.09821, 114.09821,\n",
       "          114.09821]]], dtype=float32),\n",
       " array([[[-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         ...,\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1]],\n",
       " \n",
       "        [[-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         ...,\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1]],\n",
       " \n",
       "        [[-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         ...,\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1]],\n",
       " \n",
       "        [[-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         ...,\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1]],\n",
       " \n",
       "        [[-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         ...,\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1],\n",
       "         [-1, -1, -1, ..., -1, -1, -1]]]),\n",
       " array([[[[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]]],\n",
       " \n",
       " \n",
       "        [[[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]]],\n",
       " \n",
       " \n",
       "        [[[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]]],\n",
       " \n",
       " \n",
       "        [[[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "         [[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]]]], dtype=float32),\n",
       " array([[[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "        [[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "        [[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "        [[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]],\n",
       " \n",
       "        [[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]]], dtype=float32),\n",
       " array([131.9128 , 146.68463, 175.14467, 247.13646], dtype=float32),\n",
       " array([[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "          0., 0., 0., 0., 0., 0., 0., 0., 0.]]], dtype=float32),\n",
       " {'hps': array([-37., -48., -37., -48., -37., -48., -37., -48., -37., -48., -37.,\n",
       "         -48., -37., -48., -37., -48., -37., -48., -37., -48., -37., -48.,\n",
       "         -37., -48., -37., -48., -37., -48., -37., -48., -37., -48., -37.,\n",
       "         -48.], dtype=float32),\n",
       "  'hm_hp': array([[[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]],\n",
       "  \n",
       "         [[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]],\n",
       "  \n",
       "         [[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]],\n",
       "  \n",
       "         [[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]],\n",
       "  \n",
       "         [[1.        , 0.8007374 , 0.41111228, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.8007374 , 0.6411804 , 0.329193  , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.41111228, 0.329193  , 0.16901332, ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          ...,\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ],\n",
       "          [0.        , 0.        , 0.        , ..., 0.        ,\n",
       "           0.        , 0.        ]]], dtype=float32),\n",
       "  'hp_offset': array([[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]], dtype=float32),\n",
       "  'hp_ind': array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),\n",
       "  'hp_mask': array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),\n",
       "  'ind': 3061})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_set[0]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
